Feat/add lora for sglangjax #826

aolemila · 2025-12-02T09:55:34Z

Resolves #825.

Add codes in scripts/grpo_demo_llama3_qwen2.py to run LoRA.
Add the sglang_jax_lora_test.py to ensure update_params works, and put it into tpu-tests.yml. verify_update_params will be executed when VERIFY_UPDATE_PARAMS_KEY is configured.
Add more fields for SGLangJax in RolloutConfig.
Pass the following tests. Environment: TPU-v6e-4.

Test1: Run verification of update_params

JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 tests/generate/sglang_jax_lora_test.py.

Test2: Run scripts/grpo_demo_llama3_qwen2.py without LoRA

JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 scripts/grpo_demo_llama3_qwen2.py --num-batches 2 --num-test-batches 1 --root-dir=/home/gcpuser/aolemila --rollout-engine sglang_jax.

Test3: Run scripts/grpo_demo_llama3_qwen2.py with LoRA

JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 scripts/grpo_demo_llama3_qwen2.py --num-batches 2 --num-test-batches 1 --root-dir=/home/gcpuser/aolemila --rollout-engine sglang_jax --enable-lora --lora-target-modules all.

Reference

Colab Notebook

Checklist

I have added all the necessary unit tests for my change.
I have verified that my change does not break existing code and all unit tests pass.
I have added all appropriate doc-strings/documentation.
My PR is based on the latest changes of the main branch (if unsure, rebase the code).
I have signed the Contributor License Agreement.
I have followed Contribution Guidelines.

wang2yn84 · 2025-12-19T00:57:44Z

Hi @aolemila , thank you for the PR! Can you rebase to head and resolve the conflicts? We've removed the sglang script so it should merge into our main script. And please squash the commits.

wang2yn84

Thank you for your PR! Left some comments.

scripts/grpo_demo_sglang_jax_rollout.py

tests/generate/sglang_jax_lora_test.py

tunix/generate/utils.py

aolemila · 2025-12-25T02:55:02Z

Hi @wang2yn84 , thanks for your reply. I will rebase the main and modify codes according to your advice.

aolemila · 2025-12-25T12:43:14Z

I am rerunning scripts and fix new problems I meet.

aolemila · 2025-12-26T05:11:15Z

Hi, @wang2yn84 . I have updated codes according to your suggestions. In addition to modifications, I have passed three test cases. You can see more details in PR descriptions.

Test1: Run verification of update_params
Test2: Run scripts/grpo_demo_llama3_qwen2.py without LoRA
Test3: Run scripts/grpo_demo_llama3_qwen2.py with LoRA

tunix/models/llama3/mapping_sglang_jax.py

wang2yn84 · 2026-01-15T00:39:17Z

tunix/generate/sglang_jax_sampler.py

    new_model_state_leaves, _ = jax.tree_util.tree_flatten(new_state)
    self._model_runner.model_state_leaves = new_model_state_leaves

+    flatten_src_to_tgt_module_name = os.getenv(VERIFY_UPDATE_PARAMS_KEY, None)


This part of validation should belong to the test instead of the production code?

Fixed here: commit.

wang2yn84 · 2026-01-15T00:42:01Z

tests/generate/sglang_jax_lora_test.py

@@ -0,0 +1,371 @@
+"""


Sorry I didn't look into the details in the last round of review. Seems this integration test is quite heavy, using 3B model to run the whole GRPO workflow. Such test better go to nightly regression. In CI, can we have some lightweight validation such just test update_param?

Ok. I add python scripts/grpo_demo_llama3_qwen2.py --num-batches 2 --num-test-batches 1 --root-dir=/home/gcpuser/aolemila --rollout-engine sglang_jax --enable-lora --lora-target-modules all in tpu-nightly-regression.yml to run LoRA case. For tests/generate/sglang_jax_lora_test.py in tpu-tests.yml, I will simplify it and make it more lightweight.

Fixed here: commit.

wang2yn84 · 2026-01-15T00:53:19Z

tests/generate/sglang_jax_lora_test.py

+  return text.split("####")[1].strip()
+
+
+def download_kaggle_dataset(target_dir="./data/gsm8k"):


Can we leverage the existing API? We have dataset loading and get lora model APIs. No need to recreate these functions again. If the existing API is not sufficient, say there is no other dataset support, can you help improve the existing API maybe in a separate PR?

Ok. These codes are based on old version scripts/grpo_demo_llama3_qwen2.py, and maybe they are outdated. I will follow the latest scripts/grpo_demo_llama3_qwen2.py to use recommended APIs.

This is not used in simpified version. Fixed here: commit.

wang2yn84 · 2026-01-15T00:58:02Z

tunix/rl/rollout/base_rollout.py

-  # List of batch sizes buckets for jax jit
-  rollout_sglang_jax_precompile_bs_paddings: Optional[List] = None
+  # Whether to use lora
+  rollout_sglang_jax_enable_static_lora: bool = False


Is this suppose to be True? IIUC, the way Tunix uses Lora is static, cuz we don't require to select from multiple lora and change on the fly.

There is another case that you may not use LoRA, so setting it to True ensures that you know you are using LoRA. And SGLangJax will replace the base_layer and initialize the zero buffer if you enable_static_lora. There are a few differences compared with disabling static lora.

wang2yn84 · 2026-01-15T01:03:42Z

tunix/generate/sglang_jax_sampler.py

+  if (
+      mappings is None
+      or not enable_static_lora
+      or lora_target_modules is None


"not lora_target_modules" should have the same effect as "or lora_target_modules is None or len(lora_target_modules) == 0"

wang2yn84 · 2026-01-15T01:05:26Z

tunix/generate/sglang_jax_sampler.py

    self.engine = Engine(**self.args)

-    self.mappings = config.mapping_config.to_hf_mappings
+    self.to_hf_key_mappings = config.mapping_config.to_hf_mappings


wang2yn84 · 2026-01-15T01:07:49Z

tunix/generate/sglang_jax_sampler.py

-      args["enable_deterministic_sampling"] = True
    if config.init_with_random_weights:
      args["load_format"] = "dummy"
+    args["disable_radix_cache"] = config.disable_radix_cache


Consider put checkers into a separate function

wang2yn84 · 2026-01-15T01:09:46Z

tunix/generate/sglang_jax_sampler.py

    )
+
+
+def update_hf_key_mappings_with_lora(


Probably need to move this function to top of the file, other our internal might complain about not be able to find the function.

aolemila · 2026-01-15T09:11:09Z

I run python scripts/grpo_demo_llama3_qwen2.py --num-batches 2 --num-test-batches 1 --root-dir=/home/gcpuser/aolemila --rollout-engine sglang_jax --enable-lora --lora-target-modules all.

aolemila requested review from abheesht17, hgao327, jiangyangmu, lc5211, sizhit2, tianshub and wang2yn84 as code owners December 2, 2025 09:55

aolemila had a problem deploying to testing December 2, 2025 09:55 — with GitHub Actions Failure

aolemila force-pushed the feat/add-lora-for-sglangjax branch from 32a06fd to 3af36c3 Compare December 3, 2025 09:54

aolemila had a problem deploying to testing December 3, 2025 09:54 — with GitHub Actions Failure

aolemila changed the title ~~[WIP] Feat/add lora for sglangjax~~ [Feature] Feat/add lora for sglangjax Dec 3, 2025

aolemila had a problem deploying to testing December 3, 2025 10:32 — with GitHub Actions Failure

aolemila force-pushed the feat/add-lora-for-sglangjax branch from 3ea01a4 to a56ef3a Compare December 4, 2025 06:46

aolemila had a problem deploying to testing December 4, 2025 06:46 — with GitHub Actions Error

aolemila had a problem deploying to testing December 4, 2025 06:46 — with GitHub Actions Failure

aolemila had a problem deploying to testing December 4, 2025 06:51 — with GitHub Actions Failure

aolemila force-pushed the feat/add-lora-for-sglangjax branch from 3dfcb05 to bb7bd21 Compare December 5, 2025 02:46

aolemila had a problem deploying to testing December 5, 2025 02:46 — with GitHub Actions Failure

aolemila mentioned this pull request Dec 5, 2025

[WIP] add lora codes for SglangJaxRollout #844

Closed

6 tasks

aolemila mentioned this pull request Dec 17, 2025

[Feature] Ensure to pass e2e in SGLangJaxRollout with LoRA for Tunix sgl-project/sglang-jax#417

Open

2 tasks

wang2yn84 reviewed Dec 19, 2025

View reviewed changes

aolemila force-pushed the feat/add-lora-for-sglangjax branch from bb7bd21 to e32615c Compare December 25, 2025 12:44

aolemila had a problem deploying to testing December 25, 2025 12:44 — with GitHub Actions Failure

aolemila changed the title ~~[Feature] Feat/add lora for sglangjax~~ [WIP] Feat/add lora for sglangjax Dec 25, 2025

aolemila force-pushed the feat/add-lora-for-sglangjax branch from e32615c to 3f755c5 Compare December 26, 2025 03:25

aolemila had a problem deploying to testing December 26, 2025 03:25 — with GitHub Actions Failure

aolemila temporarily deployed to testing December 26, 2025 03:25 — with GitHub Actions Inactive

aolemila changed the title ~~[WIP] Feat/add lora for sglangjax~~ Feat/add lora for sglangjax Dec 26, 2025

aolemila self-assigned this Dec 26, 2025

aolemila force-pushed the feat/add-lora-for-sglangjax branch from 3f755c5 to 874bfbc Compare December 26, 2025 04:26

aolemila had a problem deploying to testing December 26, 2025 04:26 — with GitHub Actions Failure

aolemila temporarily deployed to testing December 26, 2025 04:26 — with GitHub Actions Inactive

aolemila temporarily deployed to testing December 26, 2025 04:44 — with GitHub Actions Inactive

JamesBrianD reviewed Jan 4, 2026

View reviewed changes

tunix/models/llama3/mapping_sglang_jax.py Outdated Show resolved Hide resolved

aolemila force-pushed the feat/add-lora-for-sglangjax branch from 874bfbc to fd80b54 Compare January 7, 2026 11:30

aolemila temporarily deployed to testing January 7, 2026 11:30 — with GitHub Actions Inactive

wang2yn84 reviewed Jan 15, 2026

View reviewed changes

complete LoRA codes for SGLangJax, and pass tests

0186dd8

aolemila force-pushed the feat/add-lora-for-sglangjax branch from fd80b54 to 401c62f Compare January 15, 2026 08:59

aolemila had a problem deploying to testing January 15, 2026 09:00 — with GitHub Actions Error

update codes according to code reviews

48c01da

aolemila force-pushed the feat/add-lora-for-sglangjax branch from 401c62f to 48c01da Compare January 15, 2026 09:01

aolemila temporarily deployed to testing January 15, 2026 09:02 — with GitHub Actions Inactive

		return text.split("####")[1].strip()


		def download_kaggle_dataset(target_dir="./data/gsm8k"):

Feat/add lora for sglangjax #826

Are you sure you want to change the base?

Feat/add lora for sglangjax #826

Uh oh!

Conversation

aolemila commented Dec 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wang2yn84 commented Dec 19, 2025

Uh oh!

wang2yn84 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aolemila commented Dec 25, 2025

Uh oh!

aolemila commented Dec 25, 2025

Uh oh!

aolemila commented Dec 26, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aolemila Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aolemila commented Jan 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aolemila commented Dec 2, 2025 •

edited

Loading

aolemila Jan 15, 2026 •

edited

Loading